# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
# --------------------------------------------------------
# References:
# timm: https://github.com/rwightman/pytorch-image-models/tree/master/timm
# DeiT: https://github.com/facebookresearch/deit
# --------------------------------------------------------

from functools import partial

import torch
import torch.nn as nn

import timm.models.vision_transformer


class VisionTransformer(timm.models.vision_transformer.VisionTransformer):
    """ 带有全局平均池化支持的 Vision Transformer
    这个类扩展了 timm 库中的 VisionTransformer 模型，增加了一个全局平均池化的选项。
    如果 `global_pool` 为 True，输出特征会在经过池化后传递给最终的分类器。
    """

    def __init__(self, global_pool=False, **kwargs):
        """
        初始化 VisionTransformer 模型。

        参数：
            global_pool (bool): 是否在最终分类器前进行全局平均池化。
            **kwargs: 传递给父类 VisionTransformer 的其他参数。
        """
        # 调用父类 (VisionTransformer) 的初始化方法
        super(VisionTransformer, self).__init__(**kwargs)

        # 保存 global_pool 标志，控制是否使用全局平均池化
        self.global_pool = global_pool

        if self.global_pool:
            # 如果启用了全局池化，则为输出创建一个额外的归一化层
            norm_layer = kwargs['norm_layer']  # 从 kwargs 获取归一化层
            embed_dim = kwargs['embed_dim']  # 从 kwargs 获取嵌入维度
            self.fc_norm = norm_layer(embed_dim)  # 使用提供的 `norm_layer` 创建归一化层

            # 删除原始的归一化层 `norm`，因为它会被池化后的新归一化层替代
            del self.norm

    def forward_features(self, x):
        """
        定义 Vision Transformer 中特征提取部分的前向传播过程。

        参数：
            x (Tensor): 输入张量，形状为 (B, C, H, W)，其中 B 是批次大小，
                        C 是通道数（例如 RGB 图像是 3），H 和 W 是图像的空间维度。

        返回：
            Tensor: 经过 Vision Transformer 处理后的输出特征。
        """
        B = x.shape[0]  # 获取批次大小

        # 对输入图像进行 patch 嵌入（将图像分割为补丁并映射到嵌入空间）
        x = self.patch_embed(x)

        # 创建并将类标记（cls token）与补丁 token 拼接
        cls_tokens = self.cls_token.expand(B, -1, -1)  # 扩展 cls_token 以匹配批次大小 B
        x = torch.cat((cls_tokens, x), dim=1)  # 将类标记与补丁嵌入拼接起来

        # 添加位置嵌入到拼接后的 token（包括类标记）
        x = x + self.pos_embed

        # 对拼接后的 token 应用位置编码的 dropout
        x = self.pos_drop(x)

        # 将输入通过每个 Transformer 块（多头自注意力和前馈网络）
        for blk in self.blocks:
            x = blk(x)

        if self.global_pool:
            # 如果启用了全局池化，则对所有空间位置进行全局平均池化
            x = x[:, 1:, :].mean(dim=1)  # 排除类标记，平均池化剩余的补丁 token
            outcome = self.fc_norm(x)  # 在池化后应用归一化层
        else:
            # 如果没有启用全局池化，则对输出进行归一化，并取类标记（x[:, 0]）
            x = self.norm(x)
            outcome = x[:, 0]  # 提取类标记作为分类结果

        return outcome



# 工厂函数 vit_base_patch16、vit_large_patch16 和 vit_huge_patch14：
# 分别创建了三个不同规模的 ViT 模型，通过调用 VisionTransformer 类并传递相应的参数。
# 用于图像分类任务的
def vit_base_patch16(**kwargs):
    model = VisionTransformer(
        patch_size=16, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4, qkv_bias=True,
        norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
    return model


def vit_large_patch16(**kwargs):
    model = VisionTransformer(
        patch_size=16, embed_dim=1024, depth=24, num_heads=16, mlp_ratio=4, qkv_bias=True,
        norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
    return model


def vit_huge_patch14(**kwargs):
    model = VisionTransformer(
        patch_size=14, embed_dim=1280, depth=32, num_heads=16, mlp_ratio=4, qkv_bias=True,
        norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
    return model